

import os,glob,time
import heilsTrain
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

def Trainor(opt):
    path = r'/home/hegang/datas2/hegang/datas/meter/class_data/4_meters/'   #data path
    FT_params_path=r"/data2/enducation/train_code/predictor/save_models/mobilenet_v2-b03531040x.pth"
    save_dir=r'./save_models'
    #
    final_path = heilsTrain.retrain(path,save_dir,FT_params_path,opt)
    print("Train final_path: {} ".format(final_path))

    today=time.strftime("%Y%m%d",time.localtime())
    # save_path_list = glob.glob(os.path.join(save_dir, "20211020_mobilenet_v2_NotSigmoid_NOinit", "*.pth"))
    save_path_list = glob.glob(os.path.join(save_dir, today, "*.pth"))
    if len(save_path_list)>7:
        delete_model = sorted(save_path_list, key=(lambda x:os.path.split(x)[1][:-4].split("-")[-1]))[0:-20]
        for delete_model_name in delete_model:
            os.remove(delete_model_name)
            print("success delete :", delete_model_name)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--endepoch', type=int, default=80)
    parser.add_argument('--scheduler_lr', type=str, default='cawb', help='the learning rate scheduler, cos/cawb')
    parser.add_argument('--cawb_steps', nargs='+', type=int, default=[20,40,60],
                        help='the cawb learning rate scheduler steps')
    opt = parser.parse_args()


    Trainor(opt)
